package com.hc.app.frequencycontrol.strategy;

import com.hc.app.constant.FrequencyControlConstant;
import com.hc.app.dto.LeakBucketDTO;
import com.hc.app.frequencycontrol.AbstractFrequencyControl;
import com.hc.app.manager.LeakBucketManager;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @Author: huangcheng
 * @CreateTime: 2024-07-25  14:50
 * @Description: 漏桶算法实现
 * @Version: 1.0
 */
@Service
public class LeakBucketFrequencyControl extends AbstractFrequencyControl<LeakBucketDTO> {
    @Autowired
    private LeakBucketManager leakBucketManager;
    @Override
    protected boolean reachRateLimit(Map<String, LeakBucketDTO> frequencyControlMap) {
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            // 尝试获取 1 个令牌
            return leakBucketManager.tryAcquire(key, 1);
        }
        return false;
    }

    @Override
    protected void addFrequencyControlStatisticsCount(Map<String, LeakBucketDTO> frequencyControlMap){
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            LeakBucketDTO tokenBucketDTO = frequencyControlMap.get(key);
            leakBucketManager.createLeakBucket(key, tokenBucketDTO.getCapacity(), tokenBucketDTO.getRate());
            // 扣减 1 个令牌
            leakBucketManager.deductionToken(key, 1);
        }
    }

    @Override
    protected String getStrategyName() {
        return FrequencyControlConstant.LEAK_BUCKET;
    }
}
